/*
  S.M.A.C.K - An operating system kernel
  Copyright (C) 2010,2011 Mattias Holm and Kristian Rietveld
  For licensing and a full list of authors of the kernel, see the files
  COPYING and AUTHORS.
*/

#include <process.h>
#include <assert.h>
#include <string.h>
#include <vm.h>
#include <elf.h>
#include <interrupt.h>
#include <stdint.h>
#include <bittools.h>
#include "cpu.h"

#define PAGE_SIZE 4096
#define MIB 1048576

#define SVC_STACK_SIZE   (2*PAGE_SIZE)
#define STACK_BOTTOM (0x8000000 - (SVC_STACK_SIZE + PAGE_SIZE))
#define STACK_SIZE   (32*PAGE_SIZE)
#define STACK_TOP    (STACK_BOTTOM - STACK_SIZE)
#define ARGV_PAGE    (STACK_TOP - 2 * PAGE_SIZE)

static uint64_t addr_space_bitmap[MAX_ADDR_SPACES/64];

unsigned
arm_reserve_addr_space_id(void)
{
  int bit_id = 0;
  uint32_t flags = save_flags();
  disable_interrupts();
  for (unsigned i = 0 ; i < MAX_ADDR_SPACES/64 ; i++) {
    if (addr_space_bitmap[i] != UINT64_MAX ) {
      int bit_offs = clo_64(addr_space_bitmap[i]);
      bit_id = i * 64 + bit_offs;
      addr_space_bitmap[i] |= (1ULL << (63 - bit_offs));
      break;
    }
  }

  restore_flags(flags);
  return bit_id;
}

void
arm_release_addr_space_id(unsigned i)
{
  uint32_t flags = save_flags();
  disable_interrupts();

  addr_space_bitmap[i / 64] &= ~((1ULL << 63) >> (i % 64));

  restore_flags(flags);
}


void
hw_proc_create_user_stack(proc_t *proc)
{
  uint32_t saved_ttbr0 = read_tbr0();
  set_tbr0(proc->cpu_helper.ttbr0);

  /* Allocate space for the SYS stack and setup SYS SP register */
  va_t stack_addr = (va_t)vm_map(&proc->vm_map, VM_RW, STACK_TOP, STACK_SIZE);
  assert(stack_addr == STACK_TOP);

  proc->cpu_ctxt.gprs[13] = (STACK_BOTTOM - 4);

  set_tbr0(saved_ttbr0);
}

void
hw_proc_set_argv(proc_t *proc, const char *argv, size_t size)
{
  uint32_t saved_ttbr0 = read_tbr0();
  set_tbr0(proc->cpu_helper.ttbr0);

  if (size > PAGE_SIZE)
    size = PAGE_SIZE;

  /* Set up a page to store the argv array in for the new process */
  assert((va_t)vm_map(&proc->vm_map, VM_RW, ARGV_PAGE, PAGE_SIZE) == ARGV_PAGE);
  memcpy((void *)ARGV_PAGE, argv, size);

  set_tbr0(saved_ttbr0);
}

int
hw_proc_create(proc_t *proc)
{
  /* Set up the page tables for this process */
  void *addr = vm_map_align(vm_get_kernel_map(VM_REG_KERNEL_HEAP),
			    VM_SUPER_RW|VM_CONTIG|VM_NO_CACHE, PAGE_SIZE*2,
			    2*PAGE_SIZE + 2*MIB);
  assert(addr != NULL && "Could not allocate page table for user process");
  memset(addr, 0, 2*PAGE_SIZE + 2*MIB);

  pa_t pagetable = vm_probe_physical(vm_get_kernel_map(VM_REG_KERNEL_HEAP), addr);

  proc->cpu_helper.ttbr0 = pagetable;
  proc->cpu_helper.ttbr0_va = (va_t)addr;
  proc->vm_map.page_table = addr;
  proc->vm_map.page_table_pa = pagetable;


  /* Allocate process stack for SVC mode */
  proc->cpu_helper.svc_stack = (va_t)vm_map(vm_get_kernel_map(VM_REG_KERNEL_HEAP), VM_SUPER_RW,
                                            0, SVC_STACK_SIZE);
  assert(proc->cpu_helper.svc_stack);
  proc->cpu_ctxt.svc_sp = proc->cpu_helper.svc_stack + 2*PAGE_SIZE - 4;

  /* Allocate process stack for SYS mode */
  hw_proc_create_user_stack(proc);

  return 0;
}

int
hw_proc_load(proc_t *proc, const char *filename)
{
  uint32_t saved_ttbr0;

  saved_ttbr0 = read_tbr0();
  set_tbr0(proc->cpu_helper.ttbr0);

  /* The SPSR should start in SVC_MODE, because proc_start() will
   * run first.  In proc_start() we will make sure we enter USR_MODE
   * and jump to the user program.
   */
  proc->cpu_ctxt.spsr = 0x13;

  /* Load executable file into memory */
  proc->entrypoint = (uint32_t) elf_load(proc, filename);
  if (proc->entrypoint == 0x0)
    assert(0 && "elf_load failed");

  if (proc == proc_current())
    {
      /* There is no point in modifying the context in this case, because
       * on return from this system call, the context will not be restored
       * from memory.  We have to directly modify the SYS mode registers.
       */

      /* Set LR that will be after after SWI handling to the
       * process entrypoint.
       */
      uint32_t *swi_stack =
          (uint32_t *)(proc->cpu_helper.svc_stack
                       + 2 * PAGE_SIZE - 4 - 13 * 4);

      swi_stack[12] = proc->entrypoint;

      /* Reset USR stack, reset LR to 0x0 */
      __asm__ volatile (
                        "cps #0b11111 @ SYS_MODE\n"
                        "mov sp, %[sp_val]\n"
                        "mov lr, #0x0\n"
                        "cps #0b10011 @ SVC_MODE\n"
                        :
                        : [sp_val] "r" (STACK_BOTTOM - 4)
                       );
    }
  else
    {
      /* Set USR lr */
      proc->cpu_ctxt.gprs[14] = proc->entrypoint + 4;

      /* Set SVC lr */
      proc->cpu_ctxt.svc_lr = ((uint32_t)&hw_proc_start) + 4;
      proc->cpu_ctxt.pc = proc->cpu_ctxt.svc_lr;
    }

  set_tbr0(saved_ttbr0);

  return 0;
}

static void
transfer_region(proc_t *dup, proc_t *p, va_t addr, size_t size)
{
  /* Transfer a region from one address space to the other.
   *
   * This can be optimized by temporarily mapping in the destination
   * address range into the address space from which the data will originate.
   */
  void *data = vm_map(vm_get_kernel_map(VM_REG_KERNEL_HEAP), VM_SUPER_RW, 0,
                      size);

  uint32_t saved_ttbr0 = read_tbr0();
  set_tbr0(p->cpu_helper.ttbr0);
  memcpy(data, (void *)addr, size);
  set_tbr0(dup->cpu_helper.ttbr0);
  memcpy((void *)addr, data, size);
  set_tbr0(saved_ttbr0);

  vm_unmap(vm_get_kernel_map(VM_REG_KERNEL_HEAP), data);
}

static void
hw_proc_dup_pages(proc_t *dup, proc_t *p)
{
  vm_region_t *reg;
  uint32_t saved_ttbr0 = read_tbr0();

  /* Duplicates all pages of process p into process dup.
   *
   * We skip the stack, since this is created (from scratch) by
   * proc_create() already.
   */
  LIST_FOREACH(reg, p->vm_map.region, link)
    {
      if (reg->start == STACK_TOP)
        continue;

      set_tbr0(dup->cpu_helper.ttbr0);
      assert((va_t)vm_map(&dup->vm_map, reg->flags,
                          reg->start, reg->size) == reg->start);
      set_tbr0(saved_ttbr0);

      transfer_region(dup, p, reg->start, reg->size);
    }
}

/* NOTE: This function can only be called from system call invocation */
void
hw_proc_dup(proc_t *dup, proc_t *p)
{
  hw_proc_dup_pages(dup, p);

  /* Copy the SYS mode stack */
  transfer_region(dup, p, STACK_TOP, STACK_SIZE);

  dup->entrypoint = p->entrypoint;

  /* Register setup */
  dup->cpu_ctxt.spsr = 0x13;

  /* Set USR lr */
  dup->cpu_ctxt.gprs[14] = dup->entrypoint + 4;

  /* Set SVC lr */
  dup->cpu_ctxt.svc_lr = ((uint32_t)&hw_proc_start) + 4;
  dup->cpu_ctxt.pc = dup->cpu_ctxt.svc_lr;

  /* We cannot do a plain copy from p's CPU context, simply because
   * the context is not current.  It is only updated on context
   * switch.
   */
  __asm__ volatile (
    "cps #0b11111 @ SYS_MODE\n"
    "mov %[sp_val], sp\n"
    "mov %[lr_val], lr\n"
    "cps #0b10011 @ SVC_MODE\n"
    : [sp_val] "=r" (dup->cpu_ctxt.gprs[13]),
    [lr_val] "=r" (dup->cpu_ctxt.gprs[14])
    );

  uint32_t *swi_stack = (uint32_t *)(p->cpu_helper.svc_stack + 2 * PAGE_SIZE - 4 - 13 * 4);
  for (int i = 1; i < 13; i++)
    dup->cpu_ctxt.gprs[i] = swi_stack[i-1];

  /* Currently, there is no need for copying cpu_helper, since a
   * new page table is allocated and no other data is there.
   */

  /* We put the address where the hw_return_from_fork handler has to
   * jump to on the SVC stack.
   */
  uint32_t *dup_swi_stack = (uint32_t *)(dup->cpu_helper.svc_stack + 2 * PAGE_SIZE - 4 - 4);
  dup_swi_stack[0] = swi_stack[12];

  dup->cpu_ctxt.pc = ((uint32_t)&hw_return_from_fork) + 4;
  dup->cpu_ctxt.svc_lr = ((uint32_t)&hw_return_from_fork) + 4;

  /* Reset SVC mode stack pointer. Take address into account that was
   * pushed on the stack above.
   */
  dup->cpu_ctxt.svc_sp = dup->cpu_helper.svc_stack + 2 * PAGE_SIZE - 4 - 4;
}

void
hw_proc_unload(proc_t *proc)
{
  uint32_t saved_ttbr0;

  saved_ttbr0 = read_tbr0();
  set_tbr0(proc->cpu_helper.ttbr0);

  /* unmap kernel stack, note that the user stack will be unmapped when user
     regions are unmapped */
  vm_unmap(vm_get_kernel_map(VM_REG_KERNEL_HEAP), (void *)proc->cpu_helper.svc_stack);

  set_tbr0(saved_ttbr0);

  /* unmap page table */
  vm_unmap(vm_get_kernel_map(VM_REG_KERNEL_HEAP),
	   (void *)proc->cpu_helper.ttbr0_va);
}

void
hw_proc_preempt(proc_t *p)
{
  proc_t *current_proc = proc_current();
  extern uint32_t _irq_stack_bottom;
  uintptr_t ctxt = (uintptr_t)&_irq_stack_bottom;
  ctxt -= sizeof(cpu_ctxt_t);

  /* Store the process context stored on the IRQ stack into its proper
   * memory location.  And place another process' context on this location.
   * Also see the _isr routine in arch/arm/exception.s for more details.
   */
  memcpy(&current_proc->cpu_ctxt, (void *)ctxt, sizeof(cpu_ctxt_t));
  set_tbr0(p->cpu_helper.ttbr0);
  memcpy((void *)ctxt, &p->cpu_ctxt, sizeof(cpu_ctxt_t));
  // TODO: Change process number in MMU
}

/* FIXME: ASID code to be enabled in the future */
#if 0

// For calling after process termination
static void
arm_invalidate_tlb_by_asid(uint32_t asid)
{
    __asm__ volatile (
    "mcr p15, 0, %[asid], c8, c6, 2\n"
    : [asid] "=r" (asid)
    );
}

static void
arm_change_asid_and_ttbr0(uint32_t new_asid, uint32_t new_ttbr_val)
{
  uint32_t flags = save_flags();
  disable_interrupts();

  set_ttbcr2(TTBCR_PD0, 0);
  isb();
  // Write asid value
  __asm__ volatile (
    "mcr p15, 0, %[asid], c13, c0, 1\n"
    : [asid] "=r" (new_asid)
  );

  set_tbr0_2(new_ttbr_val);
  isb();
  set_ttbcr2(0, TTBCR_PD0);

  restore_flags(flags);
}
#endif
